#pragma once

#include "quote_data.h"
#include <string>
#include <vector>
#include <unordered_map>
#include <thread>
#include <atomic>
#include <mutex>
#include <functional>
#include <cstdint>
#include <chrono>
#include <cstring>

#ifdef _WIN32
#include <winsock2.h>
#include <ws2tcpip.h>
#else
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <unistd.h>
#endif

namespace yquote {

// WebSocket帧类型枚举
enum class WebSocketFrameType : uint8_t {
    CONTINUATION = 0x0,
    TEXT = 0x1,
    BINARY = 0x2,
    CLOSE = 0x8,
    PING = 0x9,
    PONG = 0xA,
    UNKNOWN = 0xFF
};

// WebSocket接收结果
enum class WebSocketReceiveResult {
    SUCCESS,           // 成功接收到业务消息
    PING_RECEIVED,     // 收到Ping帧
    PONG_RECEIVED,     // 收到Pong帧
    CLOSE_RECEIVED,    // 收到关闭帧
    CONNECTION_ERROR,  // 连接错误
    PROTOCOL_ERROR,    // 协议错误
    UNSUPPORTED_FRAME, // 不支持的帧类型
    NO_DATA           // 没有数据（非阻塞模式）
};

// WebSocket帧结构
struct WebSocketFrame {
    bool fin;                    // 是否为最后一个分片
    bool rsv1, rsv2, rsv3;      // 保留位
    WebSocketFrameType opcode;   // 操作码
    bool masked;                 // 是否有掩码
    uint64_t payload_length;     // 负载长度
    uint8_t mask[4];            // 掩码（如果有）
    std::vector<uint8_t> payload; // 负载数据
    
    WebSocketFrame() : fin(false), rsv1(false), rsv2(false), rsv3(false), 
                      opcode(WebSocketFrameType::UNKNOWN), masked(false), 
                      payload_length(0) {
        memset(mask, 0, sizeof(mask));
    }
};

// 客户端连接状态
enum class ClientConnectionState {
    CONNECTING,    // 正在连接
    CONNECTED,     // 已连接
    CLOSING,       // 正在关闭
    CLOSED         // 已关闭
};

class WebSocketServer {
public:
    using MessageCallback = std::function<void(const std::string& client_id, const std::string& message)>;
    using ClientConnectCallback = std::function<void(const std::string& client_id)>;
    using ClientDisconnectCallback = std::function<void(const std::string& client_id)>;
    using ErrorCallback = std::function<void(const std::string& client_id, const std::string& error)>;
    
    WebSocketServer(int port = 8080);
    ~WebSocketServer();
    
    // 启动服务器
    bool start();
    
    // 停止服务器
    void stop();
    
    // 检查是否正在运行
    bool isRunning() const;
    
    // 向指定客户端发送消息
    bool sendMessage(const std::string& client_id, const std::string& message);
    
    // 向指定客户端发送二进制数据
    bool sendBinaryMessage(const std::string& client_id, const std::vector<uint8_t>& data);
    
    // 广播消息给所有客户端
    void broadcast(const std::string& message);
    
    // 广播二进制数据给所有客户端
    void broadcastBinary(const std::vector<uint8_t>& data);
    
    // 向订阅了指定股票的客户端发送消息
    void sendToSubscribers(const std::vector<std::string>& client_ids, const std::string& message);
    
    // 发送Ping帧到指定客户端
    bool sendPing(const std::string& client_id, const std::vector<uint8_t>& payload = {});
    
    // 优雅关闭客户端连接
    bool closeClient(const std::string& client_id, uint16_t code = 1000, const std::string& reason = "");
    
    // 设置回调函数
    void setMessageCallback(MessageCallback callback);
    void setClientConnectCallback(ClientConnectCallback callback);
    void setClientDisconnectCallback(ClientDisconnectCallback callback);
    void setErrorCallback(ErrorCallback callback);
    
    // 获取连接的客户端数量
    size_t getClientCount() const;
    
    // 获取所有客户端ID
    std::vector<std::string> getClientIds() const;
    
    // 获取客户端连接状态
    ClientConnectionState getClientState(const std::string& client_id) const;
    
    // 设置心跳间隔（秒）
    void setHeartbeatInterval(int seconds);
    
private:
    struct ClientConnection {
        int socket;
        std::string id;
        std::thread thread;
        std::atomic<ClientConnectionState> state;
        std::chrono::steady_clock::time_point last_ping_time;
        std::chrono::steady_clock::time_point last_pong_time;
        std::mutex send_mutex; // 发送数据的互斥锁
        
        ClientConnection(int sock, const std::string& client_id)
            : socket(sock), id(client_id), state(ClientConnectionState::CONNECTING),
              last_ping_time(std::chrono::steady_clock::now()),
              last_pong_time(std::chrono::steady_clock::now()) {}
    };
    
    int port_;
    int server_socket_;
    std::atomic<bool> running_;
    std::thread accept_thread_;
    std::thread heartbeat_thread_;
    int heartbeat_interval_; // 心跳间隔（秒）
    
    mutable std::mutex clients_mutex_;
    std::unordered_map<std::string, std::unique_ptr<ClientConnection>> clients_;
    
    MessageCallback message_callback_;
    ClientConnectCallback connect_callback_;
    ClientDisconnectCallback disconnect_callback_;
    ErrorCallback error_callback_;
    
    // 服务器主循环
    void acceptLoop();
    
    // 心跳检测循环
    void heartbeatLoop();
    
    // 处理客户端连接
    void handleClient(const std::string& client_id);
    
    // 解析WebSocket握手
    bool performHandshake(int client_socket);
    
    // 计算WebSocket Accept值
    std::string calculateWebSocketAccept(const std::string& key);
    
    // 发送WebSocket帧
    bool sendWebSocketFrame(int socket, const WebSocketFrame& frame);
    bool sendWebSocketFrame(int socket, WebSocketFrameType type, const std::vector<uint8_t>& data);
    bool sendWebSocketFrame(int socket, const std::string& data); // 兼容旧接口
    
    // 接收WebSocket帧
    WebSocketReceiveResult receiveWebSocketFrame(int socket, WebSocketFrame& frame, std::string& message);
    
    // 发送Pong帧响应
    bool sendPongFrame(int socket, const std::vector<uint8_t>& payload);
    
    // 发送关闭帧
    bool sendCloseFrame(int socket, uint16_t code = 1000, const std::string& reason = "");
    
    // 解析WebSocket帧头
    bool parseFrameHeader(int socket, WebSocketFrame& frame);
    
    // 读取帧负载
    bool readFramePayload(int socket, WebSocketFrame& frame);
    
    // 生成客户端ID
    std::string generateClientId();
    
    // 移除客户端
    void removeClient(const std::string& client_id);
    
    // 初始化网络
    bool initializeNetwork();
    
    // 清理网络
    void cleanupNetwork();
    
    // 工具函数：安全的socket接收
    bool safeRecv(int socket, void* buffer, size_t size);
    
    // 工具函数：安全的socket发送
    bool safeSend(int socket, const void* buffer, size_t size);
    
    // 报告错误
    void reportError(const std::string& client_id, const std::string& error);
};

} // namespace yquote 